import base64
import os
import cv2
import glob
from matplotlib import pyplot as plt
from matplotlib import patches
from TrashDetector import TrashDetector


def imgEncodeBase64(filePath, outPath):
    with open(filePath, 'rb') as f1:
        base64_data = base64.b64encode(f1.read())
        with open(outPath, 'w', newline="") as f2:
            f2.write(base64_data.decode())



def main():

    with open("2.txt", 'rb') as f:
        base64code = f.read().decode("UTF-8")

    d = TrashDetector(".\\model")
    classes, trashNum, boxes, img = d.detectByBase64(base64code, codeUrlSafe=False, imgFormat="jpg")
    print(classes)
    print(trashNum)
    print(boxes)

    # plot show result
    fig, ax = plt.subplots(1, figsize=(10, 10))

    ax.imshow(img)
    # for i in range(trashNum):
    #     x1, y1, x2, y2 = boxes[i]
    #     w = x2 - x1
    #     h = y2 - y1
    #     label = classes[i]
    #
    #     if label == 0:  # trash1
    #         color = (0, 1, 0)
    #     else:           # trash2
    #         color = (1, 0, 0)
    #
    #     rect = patches.Rectangle((x1, y1), w, h, linewidth=3, edgecolor=color,
    #                              facecolor='none')
    #     ax.add_patch(rect)
    plt.waitforbuttonpress()


if __name__ == '__main__':

    main()
    # imgEncodeBase64("2.jpg", "2.txt")
    # imgsList = glob.glob("imgs\\test_2\\*.jpg")
    # outList = ["imgs\\imgsBase64_2\\" + i.split("\\")[-1][:-3] + "txt" for i in imgsList]
    #
    # print(imgsList)
    # print(outList)
    #
    # for i in range(imgsList.__len__()):
    #     imgEncodeBase64(imgsList[i], outList[i])
